package main

import (
	"fmt"
	sentinel "github.com/alibaba/sentinel-golang/api"
	"github.com/alibaba/sentinel-golang/core/base"
	"github.com/alibaba/sentinel-golang/core/flow"
	"go.uber.org/zap"
	"log"
	"math/rand"
	"time"
)

// QPS 基于QPS的限流
func QPS() {

	err := sentinel.InitDefault()
	if err != nil {
		log.Fatal("初始化sentinel异常: ", err)
	}

	// 配置限流规则
	_, err = flow.LoadRules([]*flow.Rule{
		{
			Resource:               "test1",
			TokenCalculateStrategy: flow.Direct,
			ControlBehavior:        flow.Reject, // 直接拒绝
			Threshold:              10,          // 10个流量
			StatIntervalInMs:       1000,        // 1秒钟
		},
		{
			Resource:               "test2",
			TokenCalculateStrategy: flow.Direct,
			ControlBehavior:        flow.Reject, // 直接拒绝
			Threshold:              10,
			StatIntervalInMs:       1000,
		},
	})
	if err != nil {
		zap.S().Fatal("配置限流规则失败: ", err)
	}

	// 流量控制的入口点
	for i := 0; i < 15; i++ {
		fmt.Println(i)
		exit, blockErr := sentinel.Entry("test1", sentinel.WithTrafficType(base.Inbound))
		if blockErr != nil {
			log.Fatal("Block error: ", blockErr)
		} else {
			fmt.Println(i)
			exit.Exit()
		}
	}
}

// WarmUp  预热冷启动，应对常态低流量，突然遇到激增流量，进行系统的逐步预热，让通过的流量缓步增加
func WarmUp() {

	err := sentinel.InitDefault()
	if err != nil {
		log.Fatal("初始化sentinel异常: ", err)
	}
	ch := make(chan struct{})
	// 配置限流规则
	_, err = flow.LoadRules([]*flow.Rule{
		{
			Resource:               "warmUp",
			TokenCalculateStrategy: flow.WarmUp,
			WarmUpPeriodSec:        30,          // 设置预热时长30s
			WarmUpColdFactor:       3,           // 预热因子，默认为3
			ControlBehavior:        flow.Reject, // 直接拒绝
			Threshold:              2000,        // 1000个流量
		},
	})
	if err != nil {
		zap.S().Fatal("配置限流规则失败: ", err)
	}

	// 每一秒统计一次，通过多少流量，block多少流量

	total := 0
	blockNum := 0
	pass := 0
	for i := 0; i < 3; i++ {

		go func() {
			for {
				total++
				exit, block := sentinel.Entry("warmUp", sentinel.WithTrafficType(base.Inbound))
				if block != nil { // 被block了
					blockNum++
					time.Sleep(time.Duration(rand.Uint64()) % 10 * time.Millisecond)
				} else {
					pass++
					exit.Exit()
					time.Sleep(time.Duration(rand.Uint64()) % 10 * time.Millisecond)
				}

			}

		}()
	}

	go func() {
		for {
			fmt.Println(total)
			fmt.Println(blockNum)
			fmt.Println(pass)
			time.Sleep(time.Second)
		}
	}()

	<-ch
}

func main() {
	WarmUp()
}
